#include "memory.h"

#define ALREADY_USED_MEMORY 0x2100000
#define ALREADY_USED_PAGE (ALREADY_USED_MEMORY >> 12)
#define ALREADY_USED_BITMAP_BYTES (ALREADY_USED_MEMORY >> 15)

STATIC DWORD dwPage;
STATIC BYTE MmBitmap[MEMORY_BITMAP_SIZE];
STATIC DWORD dwLastIndex;

void MmInitializeBitmap(void) {
  dwPage = TOTAL_MEMORY_BYTES >> 15;
  if (dwPage == 0)
    dwPage = MEMORY_BITMAP_SIZE;
  dwLastIndex = 0;
  RtlSetMemory(MmBitmap, 0xFFFFFFFF, ALREADY_USED_BITMAP_BYTES);
  RtlZeroMemory(MmBitmap + ALREADY_USED_BITMAP_BYTES,
                MEMORY_BITMAP_SIZE - ALREADY_USED_BITMAP_BYTES);
}

DWORD_PTR MmAllocatePage(void) {
  register DWORD i = dwLastIndex;
  do {
    if (MmBitmap[i] != 0xFF) {
      register DWORD j;
      dwLastIndex = i;
      for (j = 0; j < 8; j++)
        if (((MmBitmap[i] >> j) & 1) == 0) {
          MmBitmap[i] |= (1 << j);
          return ((i << 3) + j) << 12;
        }
    }
    ++i;
    if (i >= dwPage)
      i = 0;
  } while (i != dwLastIndex);
  return NULL;
}

void MmAllocatePages(BOOLEAN (*lpfnCallback)(DWORD_PTR)) {
  register DWORD i = dwLastIndex;
  do {
    if (MmBitmap[i] != 0xFF) {
      register DWORD j;
      for (j = 0; j < 8; j++)
        if (((MmBitmap[i] >> j) & 1) == 0) {
          MmBitmap[i] |= (1 << j);
          if (!lpfnCallback(((i << 3) + j) << 12)) {
            dwLastIndex = i;
            return;
          }
        }
    }
    ++i;
    if (i == dwPage) {
      i = 0;
    }
  } while (i != dwLastIndex);
  lpfnCallback(NULL);
}

void MmFreePage(DWORD_PTR dwAddress) {
  if (dwAddress & 0xFFF)
    return;
  dwAddress >>= 12;
  MmBitmap[dwAddress >> 3] ^= (1 << ((dwAddress & 0x7)));
}

void MmAccessPage(DWORD_PTR dwLinear, DWORD_PTR dwAddress) {
  if (dwLinear < 0xc0000000) {
    return;
  }
  dwLinear -= 0xc0000000;
  *(LPDWORD)(DWORD_PTR)(0x2000000 + ((dwLinear >> 12) << 2)) =
      (dwAddress & 0xFFFFF000) | 0x7;
}
